// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/init/init.h>
#include <turbo/flags/flag.h>
#include <gtest/gtest.h>
#include <unordered_set>

#include <pollux/testing/exec/aggregation_fuzzer_options.h>
#include <pollux/testing/exec/aggregation_fuzzer_runner.h>
#include <pollux/testing/exec/duck_query_runner.h>
#include <pollux/testing/exec/transform_result_verifier.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/functions/sparksql/aggregates/register.h>
#include <pollux/serializers/compact_row_serializer.h>
#include <pollux/serializers/presto_serializer.h>
#include <pollux/serializers/unsafe_row_serializer.h>

TURBO_FLAG(int64_t,
    seed,
    0,
    "Initial seed for random number generator used to reproduce previous "
    "results (0 means start with random seed).");

TURBO_FLAG(std::string,
    only,
    "",
    "If specified, Fuzzer will only choose functions from "
    "this comma separated list of function names "
    "(e.g: --only \"min\" or --only \"sum,avg\").");

int main(int argc, char** argv) {
  kumo::pollux::functions::aggregate::sparksql::registerAggregateFunctions(
      "", false);

  ::testing::InitGoogleTest(&argc, argv);

  // Calls common init functions in the necessary order, initializing
  // singletons, installing proper signal handlers for better debugging
  // experience
  melon::Init init(&argc, &argv);

  kumo::pollux::functions::prestosql::registerInternalFunctions();
  if (!isRegisteredNamedVectorSerde(
          kumo::pollux::VectorSerde::Kind::kPresto)) {
    kumo::pollux::serializer::presto::PrestoVectorSerde::
        registerNamedVectorSerde();
  }
  if (!isRegisteredNamedVectorSerde(
          kumo::pollux::VectorSerde::Kind::kCompactRow)) {
    kumo::pollux::serializer::CompactRowVectorSerde::
        registerNamedVectorSerde();
  }
  if (!isRegisteredNamedVectorSerde(
          kumo::pollux::VectorSerde::Kind::kUnsafeRow)) {
    kumo::pollux::serializer::spark::UnsafeRowVectorSerde::
        registerNamedVectorSerde();
  }
  kumo::pollux::memory::MemoryManager::initialize({});

  // TODO: List of the functions that at some point crash or fail and need to
  // be fixed before we can enable. Constant argument of bloom_filter_agg cause
  // fuzzer test fail.
  std::unordered_set<std::string> skipFunctions = {"bloom_filter_agg"};

  using kumo::pollux::exec::test::TransformResultVerifier;

  auto makeArrayVerifier = []() {
    return TransformResultVerifier::create("\"$internal$canonicalize\"({})");
  };

  // The results of the following functions depend on the order of input
  // rows. For some functions, the result can be transformed to a value that
  // doesn't depend on the order of inputs. If such transformation exists, it
  // can be specified to be used for results verification. If no transformation
  // is specified, results are not verified.
    std::unordered_map<
      std::string,
      std::shared_ptr<kumo::pollux::exec::test::ResultVerifier>>
      customVerificationFunctions = {
          {"last", nullptr},
          {"last_ignore_null", nullptr},
          {"first", nullptr},
          {"first_ignore_null", nullptr},
          {"max_by", nullptr},
          {"min_by", nullptr},
          // If multiple values have the same greatest frequency, the return
          // value is indeterminate.
          {"mode", nullptr},
          {"skewness", nullptr},
          {"kurtosis", nullptr},
          {"collect_list", makeArrayVerifier()},
          {"collect_set", makeArrayVerifier()},
          // Nested nulls are handled as values in Spark. But nested nulls
          // comparison always generates null in DuckDB.
          {"min", nullptr},
          {"max", nullptr},
      };

  size_t initialSeed = turbo::get_flag(FLAGS_seed) == 0 ? std::time(nullptr) : turbo::get_flag(FLAGS_seed);
  std::shared_ptr<kumo::pollux::memory::MemoryPool> rootPool{
      kumo::pollux::memory::memoryManager()->addRootPool()};
  auto duckQueryRunner =
      std::make_unique<kumo::pollux::exec::test::DuckQueryRunner>(
          rootPool.get());
  duckQueryRunner->disableAggregateFunctions(
      {// https://github.com/facebookincubator/pollux/issues/7677
       "max_by",
       "min_by",
       // The skewness functions of Pollux and DuckDB use different
       // algorithms.
       // https://github.com/facebookincubator/pollux/issues/4845
       "skewness",
       // Spark's kurtosis uses Pearson's formula for calculating the kurtosis
       // coefficient. Meanwhile, DuckDB employs the sample kurtosis calculation
       // formula. The results from the two methods are completely different.
       "kurtosis"});

  using Runner = kumo::pollux::exec::test::AggregationFuzzerRunner;
  using Options = kumo::pollux::exec::test::AggregationFuzzerOptions;

  Options options;
  options.onlyFunctions = turbo::get_flag(FLAGS_only);
  options.skipFunctions = skipFunctions;
  options.customVerificationFunctions = customVerificationFunctions;
  options.orderableGroupKeys = true;
  return Runner::run(initialSeed, std::move(duckQueryRunner), options);
}
